﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Corecalc;
using System.Diagnostics;

namespace Sisal
{
    public static partial class Procedures
    {
        //public static void InitCalls(this GRNode node) {
        //    List<GRNode> ns = node.GetLocalGraph();

        //    foreach (GRNode n in ns) {
        //        if (n is BuiltInFunctionCallNode) {

        //            BuiltInFunctionCallNode mfn = n as BuiltInFunctionCallNode;
        //            Func<GRNode[], double> timeEstimator = Function.TimeEstimatorFor(mfn.Name);
        //            n.ExecutionTimeCost = timeEstimator(mfn.InputEdges.ToArray());

        //        } else if (n is ParallelNode) {
        //            throw new System.NotImplementedException();
        //        } else if (n is CompoundNode) {

        //            foreach (GRNode nn in n.InputEdges) {
        //                nn.InitCalls();
        //            }

        //        }
        //    }
        //}

        //public static void AssignCosts(this GRNode g) {
        //    List<GRNode> nodes = g.GetLocalGraph();

        //    foreach (GRNode node in nodes) {
        //        if (node is ParallelNode) {
        //            throw new System.NotImplementedException("...");
        //        } else if (node is CompoundNode) {
        //            CompoundNode cn = node as CompoundNode;
        //            foreach (GRNode sg in cn.Graphs)
        //                sg.AssignCosts();

        //            cn.ExecutionTimeCost = cn.ComputeExecutionTimeCost();
        //        }
        //    }

        //    //Prevent recursive calls to non-formula nodes from accumulating costs
        //    if (g.IsFormulaNode) {
        //        foreach (GRNode node in nodes) {
        //            if (node != g)
        //                g.ExecutionTimeCost += node.ExecutionTimeCost;
        //        }
        //    }
        //}

        //public static void UnknownSarkarCostSetter(GRGraph g) {
        //    //Topologically sorted list (not reversed, due to the use of a support graph instead a dependency graph)
        //    //g.FormulaNodes is topologically sorted

        //    foreach (GRNode formulaNode in g.FormulaNodes) {
        //        formulaNode.InitCalls();
        //        formulaNode.AssignCosts();
        //    }
        //}

        public static Partition PartitionFunctions(GRGraph g, out int bIter, out int cIter) {
            //NB: The formula nodes in g must be reverse topologically sorted - or just topologically in our case (support graph vs. dependency graph)

            //Iterate over formula nodes AND ConstNode if such nodes exist at the highest level in the graph
            foreach (GRNode node in g.Nodes) {
                node.Overhead = MultiProcessorModel.TSched +
                                       MultiProcessorModel.bytesToCycles(node.T_in) +
                                       MultiProcessorModel.bytesToCycles(node.T_out);

                node.CP = node.ExecutionTimeCost +
                                           MultiProcessorModel.bytesToCycles(node.Overhead);

                //Assign costs and CP to root to be read in PartitionGraph
                g.Root.ExecutionTimeCost += node == g.Root ? 0.0 : node.ExecutionTimeCost;
                g.Root.CP += node == g.Root ? 0.0 : node.CP;
            }

            List<GRNode> test_pn = new List<GRNode>();
            Dictionary<GRNode, Task> test_tp = new Dictionary<GRNode, Task>();

            List<GRNode> pn = new List<GRNode>();
            Dictionary<GRNode, Task> tp = new Dictionary<GRNode, Task>();


            foreach (GRNode node in g.Nodes) {
                //Test
                Task test_task = new Task();
                test_task.AddNode(node, true);
                test_tp.Add(node, test_task);
                SetTaskCosts(test_task);

                //Actual
                Task task = new Task();
                task.AddNode(node, true);
                tp.Add(node, task);
                SetTaskCosts(task);

                test_pn.Add(node);
                pn.Add(node);
            }

            Partition test_partition = new Partition(test_pn, test_tp);
            Tuple<Partition, int, int> bestIter = PartitionGraph(g, test_partition);

            //if bestIter > (curIter - 1) then dont waste time on reconstruction
            if (bestIter.Item2 > (bestIter.Item3 - 1)) {
                bIter = bestIter.Item2;
                cIter = bestIter.Item3;
                return bestIter.Item1;
            }

            Partition partition = new Partition(pn, tp);
            Tuple<Partition, int, int> bestPartition = PartitionGraph(g, partition, true, bestIter.Item2);

            bIter = bestIter.Item2;
            cIter = bestIter.Item3;

            return bestPartition.Item1;
        }

        //public static List<GRNode> InitPrimitiveNodes(GRNode g) {
        //    List<GRNode> ns = g.GetLocalGraph();
        //    List<GRNode> pn = new List<GRNode>();

        //    foreach (GRNode n in ns) {
        //        if (n.ExecutionTimeCost < MultiProcessorModel.TMin) {
        //            pn = pn.Union(new List<GRNode> { n }).ToList();

        //            if (n is BuiltInFunctionCallNode) {
        //                BuiltInFunctionCallNode bfn = n as BuiltInFunctionCallNode;
        //                bfn.Par = false;
        //            } else if (n is ParallelNode) {
        //                ParallelNode ppn = n as ParallelNode;
        //                ppn.Size = (int)Math.Ceiling(ppn.F);
        //            }
        //        } else if (n is SimpleNode) {
        //            pn = pn.Union(new List<GRNode> { n }).ToList();
        //        } else if (n is BuiltInFunctionCallNode) {
        //            BuiltInFunctionCallNode bfn = n as BuiltInFunctionCallNode;
        //            pn = pn.Union(new List<GRNode> { n }).ToList();
        //            bfn.Par = true;
        //        } else if (n is ParallelNode) {
        //            //ParallelNode ppn = n as ParallelNode;
        //            //InitPrimitiveNodes(n);
        //            //ppn.Size = 1;
        //            throw new System.NotImplementedException();
        //        } else {
        //            //Compound Node
        //            CompoundNode node = n as CompoundNode;
        //            foreach (GRNode n2 in node.Graphs) {
        //                List<GRNode> pn2 = InitPrimitiveNodes(n2);
        //                pn = pn.Union(pn2).ToList();
        //            }
        //            pn = pn.Union(new List<GRNode> { node }).ToList();
        //        }
        //    }
            
        //    return pn;
        //}

        //Sets the cost of a task after an ACTUAL merge
        public static void SetTaskCosts(Task t) {
            t.O = MultiProcessorModel.TSched + t.T_in() + t.T_out();
            t.CP = t.T() + t.O;
        }

        //Sets task cost of tasks originated from testing merge candidates
        public static void SetTempTaskCosts(Task t) {
            t.O = MultiProcessorModel.TSched + t.T_in_Temp() + t.T_out_Temp();
            t.CP = t.T() + t.O;
        }

        public static Tuple<Partition, int, int> PartitionGraph(GRGraph g, Partition p, bool reconstruct = false, int bestIndex = 0) {
            int bestIter = 0; 
            int curIter = 0; 
            double bestCost = p.CostFunction(g.Root.ExecutionTimeCost);
            double overheadSum = p.Tasks.Sum(t => t.Q * t.O);
            double oldOverhead = 0.0;
            Task T_new = null;
            List<double> costs = new List<double>();
            List<double> cps = new List<double>();
            List<double> overheads = new List<double>();

            while (p.Tasks.Distinct().Count() > 1) {

                var tasks = p.Tasks.Distinct();

                if (reconstruct && curIter == bestIndex)
                    break;

                curIter = curIter + 1;

                var T_a = p.Tasks
                           .Distinct()
                           .OrderBy(t => t.Q * t.O).Last();


                Dictionary<GRNode, Task> localtp;
                Dictionary<Task, HashSet<Task>> path;

                SetupMerge(T_a, p, out localtp, out path);

                List<Task> sameParentAsT_a = p.TP.Values.Where(t => t != T_a).Distinct().ToList();
                Task T_b = null;
                double doesntMatter;
                double bestCP = double.MaxValue;

                List<Task> resultTasks = new List<Task>();

                foreach (var T_b_candidate in sameParentAsT_a) {
                    Task resultTask;
                    MergeTasks(T_a, T_b_candidate, localtp, false, path, out resultTask, out doesntMatter);
                    SetTempTaskCosts(resultTask);
                    SetCPUntilRoot(g, resultTask, p);

                    resultTasks.Add(resultTask);

                    if (resultTask.CP < bestCP) {
                        bestCP = resultTask.CP;
                        T_b = T_b_candidate;
                    }
                }

                //Merge T_a with the best candidate for T_b
                MergeTasks(T_a, T_b, localtp, true, path, out T_new, out oldOverhead);
                p.MergeTasks(T_a, T_b, T_new);

                SetTaskCosts(T_new);
                SetCPUntilRoot(g, T_new, p);

                overheadSum = overheadSum - oldOverhead + T_new.Q * T_new.O;
                int processors = MultiProcessorModel.ProcessorCount;

                var overheadTerm = 1 + overheadSum / g.Root.ExecutionTimeCost;
                var criticalPathTerm = g.Root.CP * processors / g.Root.ExecutionTimeCost;

                string biggest = "";
                if (criticalPathTerm > overheadTerm) {
                    p.Cost = criticalPathTerm;
                    biggest = "Critical Path Term";
                } else if (criticalPathTerm < overheadTerm) {
                    p.Cost = overheadTerm;
                    biggest = "Overhead Term";
                } else {
                    p.Cost = criticalPathTerm;
                    biggest = "Two terms are equally big";
                }

                if (!reconstruct) {
                    Console.WriteLine("CurIter: {0}, BestIter: {1}, Tasks: {2} \t\t O = 1 + overheadSum [{3:0.00}] / g.Root.Ft [{4:0.00}] = {5:0.00} \t\t CP = g.Root.CP [{6:0.00}] * P [{7}] / g.Root.Ft [{8:0.00}] = {9:0.00} \t\t\t Biggest: {10}", 
                                                              curIter, 
                                                              bestIter, 
                                                              tasks.Count(), 
                                                              1 + overheadSum, 
                                                              g.Root.ExecutionTimeCost, 
                                                              overheadTerm, 
                                                              g.Root.CP, 
                                                              processors, 
                                                              g.Root.ExecutionTimeCost, 
                                                              criticalPathTerm,
                                                              biggest);
                }

                //p.Cost = Math.Max(criticalPathTerm, overheadTerm);

                costs.Add(p.Cost);
                cps.Add(g.Root.CP);
                overheads.Add(overheadSum);

                if (p.Cost < bestCost) {
                    bestIter = curIter;
                    bestCost = p.Cost;
                }
            }

            if (reconstruct) {
                return new Tuple<Partition, int, int>(p, curIter, curIter);
            } else {
                return new Tuple<Partition, int, int>(p, bestIter, curIter);
            }
        }

        private static void SetCPUntilRoot(GRGraph g, Task T_new, Partition p) {
            GRGraph gPrime = T_new.Parent;
            while (true) {
                double oldGPrimeCP = gPrime.Root.CP;

                SetGraphCP(gPrime, p);
                if (gPrime == g || (gPrime.Root.CP == oldGPrimeCP))
                    break;

                throw new Exception("Uh oh..");
                //GRNode nPrime = gPrime.Parent;

                //double oldNPrimeCP = nPrime.CP;
                //SetENCP(nPrime);
                //gPrime = nPrime.Parent;
                //if (nPrime.CP == oldNPrimeCP)
                //    break;
            }
        }

        public static void SetupMerge(Task T_a, Partition p, out Dictionary<GRNode, Task> ltp, out Dictionary<Task, HashSet<Task>> pth) {
            Dictionary<GRNode, Task> localtp = new Dictionary<GRNode, Task>();
            Action<GRNode, Task> addTaskToLocalTp = (sn, t) => {
                sn.TempTask = t;
                sn.Task = t;
                localtp.Add(sn, t);
            };
 
            GRGraph g = T_a.Parent;
            foreach (GRNode n in g.Nodes) {
                if (p.PN.Contains(n)) {
                    var c = p.TP[n].Nodes.Count();
                    addTaskToLocalTp(n, p.TP[n]);
                } else {
                    Task t = new Task();
                    t.AddNode(n, true);

                    addTaskToLocalTp(n, t);
                }
            }

            Dictionary<Task, HashSet<Task>> path = new Dictionary<Task, HashSet<Task>>();
            foreach (Task t_i in localtp.Values.Distinct()) {
                HashSet<Task> reachable = new HashSet<Task>();

                foreach (GRNode node in t_i.Nodes) {
                    node.DFS(n => reachable.Add(n.Task));
                }

                path.Add(t_i, reachable);
            }

            ltp = localtp;
            pth = path;
        }

        public static void MergeTasks(Task t_a, Task t_b, Dictionary<GRNode, Task> localtp, bool final, Dictionary<Task, HashSet<Task>> path, out Task t, out double overhead) {
            Debug.Assert(t_a.Parent == t_b.Parent);

            Task t_new = new Task();
            double oldOverhead = 0.0;

            foreach (Task t_i in localtp.Values.Distinct()) {
                bool b1 = path[t_a].Contains(t_i), b2 = path[t_b].Contains(t_i), b3 = path[t_i].Contains(t_a), b4 = path[t_i].Contains(t_b);

                if ((b1 || b2) && (b3 || b4)) {
                    foreach (GRNode node in t_i.Nodes) {
                        t_new.AddNode(node, final);
                    }

                    oldOverhead = oldOverhead + t_i.Q * t_i.O;
                }
            }

            t = t_new;
            overhead = oldOverhead;
        }

        public static void SetGraphCP(GRGraph g, Partition p) {
            //Compute new localtp - different from the localtp in SetupMerge and MergeTasks
            Dictionary<GRNode, Task> localtp = new Dictionary<GRNode, Task>();
            Action<GRNode, Task> addTaskToLocalTp = (sn, t) => {
                sn.TempTask = t;
                sn.Task = t;
                localtp.Add(sn, t);
            };

            foreach (GRNode n in g.Nodes) {
                if (p.PN.Contains(n)) {
                    addTaskToLocalTp(n, p.TP[n]);
                } else {
                    Task t = new Task();
                    t.AddNode(n, true);

                    addTaskToLocalTp(n, t);
                }
            }

            Dictionary<Task, HashSet<Task>> prec = new Dictionary<Task, HashSet<Task>>();
            foreach (GRNode n_a in g.Nodes) {
                foreach (GRNode n_b in g.Nodes) {
                    if (!prec.ContainsKey(n_a.Task))
                        prec.Add(n_a.Task, new HashSet<Task>());

                    if (localtp[n_a] != localtp[n_b] && n_b.InputEdges.Contains(n_a)) {
                        prec[n_a.Task].Add(n_b.Task);
                    }
                }
            }

            List<Task> list = Procedures.TopSort(prec.Keys, t => prec[t]);

            //Order of the task in list "list"
            //Is it OK to assume that any task's start must be equal to the sum of "start" and "cp" of the
            //start lying directly behind it in the list?
            for (int j = 0; j < list.Count(); j++) {
                int i = j - 1;
                list[j].Start = Math.Max(0, j == 0 ? 0 : list[i].Start + list[i].CP);
            }

            g.Root.CP = Math.Max(0, list.Max(t => t.Start + t.CP));
        }

        //public static void SetENCP(GRNode n) {
        //    if (n is ParallelNode) {
        //        (n as ParallelNode).CP = (n as ParallelNode).Graph.CP;
        //    } else if (n is CompoundNode) {
        //        CompoundNode cn = n as CompoundNode;
        //        cn.CP = cn.Graphs.Sum(g => g.F * g.CP); //Does g.CriticalPath represent the sum of the CP values of the entire graph or only its own?
        //    } else {
        //        var a = "?";
        //    }
        //}

        public static IEnumerable<GRNode> TopSort(Dictionary<GRNode, List<GRNode>> outgoing) {
            return Procedures.TopSort(outgoing.Keys, n => outgoing[n]);
        }

        public static List<A> TopSort<A>(IEnumerable<A> items, Func<A, IEnumerable<A>> outgoing) {
            List<A> sortedList = new List<A>();
            List<A> marked = new List<A>();
            List<A> tempMarked = new List<A>();

            Action<A> visit = null;
            visit = n => {
                if (tempMarked.Contains(n)) throw new Exception("Uh oh..!!");
                if (!marked.Contains(n)) {
                    tempMarked.Add(n);

                    var o = outgoing(n);
                    foreach (A m in o)
                        visit(m);

                    marked.Add(n);
                    tempMarked.Remove(n);
                    sortedList.Add(n);      //Should've been AddFirst, we reverse the list instead
                }
            };

            Queue<A> unmarked = new Queue<A>(items);
            while (unmarked.Count > 0) {
                A n = unmarked.Dequeue();
                visit(n);
            }

            sortedList.Reverse();
            return sortedList;
        }
    }

    public class Partition
    {
        public List<GRNode> PN { get; set; }
        public Dictionary<GRNode, Task> TP { get; set; }

        public double Cost { get; set; }

        public Partition(List<GRNode> pn, Dictionary<GRNode, Task> tp) {
            PN = pn;
            TP = tp;
        }

        public double T_total() {
            return TP.Values.Distinct().Sum(t => t.Q * (t.T() + t.O));
        }

        public double T_crit() {
            return PN.Sum(n => n.ExecutionTimeCost) + Tasks.Distinct().Sum(t => t.O);
        }

        public void MergeTasks(Task t_a, Task t_b, Task t_new) {
            foreach (var node in t_new.Nodes) {
                TP[node] = t_new;
            }

            int taCount = t_a.Nodes.Count();
            int tbCount = t_b.Nodes.Count();
            int tNewCount = t_new.Nodes.Count();
        }

        public void MergeComplexNodeSubGraphs(Task t_new, List<GRNode> subGraphNodes) {
            foreach (var key in subGraphNodes) {
                //TP[key] = t_new;
                TP.Remove(key);
            }
        }

        public double CostFunction(double T_seq) {
            if (Cost == 0.0) {
                double P = MultiProcessorModel.ProcessorCount;

                return Math.Max(T_crit(), T_total() / P) / (T_seq / P);
            } else {
                return Cost;
            }
        }

        public IEnumerable<Task> Tasks {
            get { return TP.Values; }
        }
    }

    public class Task
    {
        public double O { get; set; }
        public double CP { get; set; }
        public double Start { get; set; }

        public List<GRNode> Nodes { get; set; }

        public Task() {
            Nodes = new List<GRNode>();
        }

        public GRNode First {
            get { return Nodes[0]; }
        }

        public bool IsSingleton {
            get { return Nodes.Count == 1; }
        }

        //Q is one in a spreadsheet context
        public double Q {
            get { return 1.0; }
        }

        public double T() {
            return Nodes.Sum(n => n.ExecutionTimeCost);
        }

        public GRGraph Parent {
            //All nodes in a task have the same parent
            get { return Nodes.Count() == 0 ? null : Nodes[0].Parent; }
        }

        public bool IncludesAllNodesInGraph(GRGraph g) {
            var exceptNodes = g.Nodes.Except(Nodes);
            return exceptNodes.Count() == 0;
        }

        public double T_in() {
            return Nodes.Sum(n => {
                return n.InputEdges.Sum(ie => {
                    return ie.Task == this ? 0.0 : MultiProcessorModel.Rc(ie.T_out);
                });
            });
        }

        public double T_out() {
            return Nodes.Sum(n => {
                return n.OutputEdges.Sum(oe => {
                    return oe.Task == this ? 0.0 : MultiProcessorModel.Wc(n.T_out);
                });
            });
        }

        public double T_in_Temp() {
            return Nodes.Sum(n => {
                return n.InputEdges.Sum(ie => {
                    return ie.Task == this ? 0.0 : MultiProcessorModel.Rc(ie.T_out);
                });
            });
        }

        public double T_out_Temp() {
            return Nodes.Sum(n => {
                return n.OutputEdges.Sum(oe => {
                    return oe.TempTask == this ? 0.0 : MultiProcessorModel.Wc(n.T_out);
                });
            });
        }

        public void AddNode(GRNode n, bool actualMerge) {
            if (actualMerge) {
                n.Task = this;

                var parents = Nodes.Select(nn => nn.Parent).Distinct();
                if (parents.Count() > 0 && !parents.Contains(n.Parent)) {
                    throw new Exception("");
                }
            }

            n.TempTask = this;
            Nodes.Add(n);
        }

        public GRNode GetNode(int i) {
            return Nodes[i];
        }

        public A FoldNodes<A>(A init, Func<A, GRNode, A> f) {
            A a = init;
            foreach (GRNode node in Nodes) {
                a = f(a, node);
            }
            return a;
        }

        public override string ToString() {
            return string.Format("Task [HashCode: {0}, Nodes: {1}, Start: {2}, O: {3}, CP: {4}]", this.GetHashCode(), Nodes.Count(), Start, O, CP);
        }
    }
}
